#!/usr/bin/python

"""Automate throughput testing across two hosts using iperf.
Dependencies:
    ssh.py              (http://commandline.org.uk/code/#sshpy)
        paramiko        (http://www.lag.net/paramiko/)
            pycrypto    (http://www.amk.ca/python/code/crypto.html)
    iperf
"""

import csv
import datetime
import ssh
import subprocess
import threading
import time
from optparse import OptionParser

WAIT_TIME = 600

def main():
    usage = "usage: %prog host [options]"
    parser = OptionParser(usage=usage)
    parser.add_option("-k", "--keyfile",
                      type="string", dest="private_key",
		              default=None,
                      help="ssh private key file")
    parser.add_option("-u", "--username",
                      type="string", dest="username",
                      help="username to log in as")
    parser.add_option("-p", "--password",
                      action="store_true", dest="password",
                      default=False, 
                      help="prompt for ssh password")
    (options, args) = parser.parse_args()

    if len(args) == 1:
        if not options.private_key and not options.password:
            print "requires ssh key or password"
            return
        while not options.username:
            options.username = raw_input("username: ")
        if options.password:
            import getpass
            password = getpass.getpass("password: ")
        else:
            password = None
    else:
        parser.print_usage() 
        return
    
    while 1:
        myssh = ssh.Connection(host=args[0], username=options.username,
                               private_key=options.private_key,
                               password=password)
        # new thread since ssh command execution will block
        print "starting remote iperf server at %s" % args[0]
        t = threading.Thread(target=run_remote_iperf, args=(myssh,))
        t.start()
        time.sleep(1)

        print "starting transfer from local iperf client"
        local_cmd = 'iperf -c {0}'.format(args[0]).split()
        p = subprocess.Popen(local_cmd, stdout=subprocess.PIPE,
                             stderr=subprocess.STDOUT)
        p.wait()
        print "transfer complete"

        print "killing remote iperf server"
        myssh.execute('kill -9 `ps aux | grep iperf | grep -v grep'
                      '| tr -s \' \' | cut -d \' \' -f2`')
        t.join()
        myssh.close()

        row = []
        row.append(datetime.datetime.now().strftime('%m/%d/%y %H:%M:%S'))
        row.append("".join(p.stdout.readlines()[-1].split()[-2:]))

        print "logging to stats.csv"
        with open('stats.csv', 'a+') as f:
            stats_csv = csv.writer(f, delimiter=',')
            stats_csv.writerow(row)

        print "waiting (%d minutes)" % (WAIT_TIME/60)
        time.sleep(WAIT_TIME)

def run_remote_iperf(sshconn):
    sshconn.execute('iperf -s')


if __name__ == '__main__':
    main()
